package cn.com.pc.content.service.es.impl;

import cn.com.pc.content.dto.PageResult;
import cn.com.pc.content.service.es.ESQueryService;
import cn.com.pc.content.service.es.ElasticEntity;
import cn.com.pc.content.util.JsonUtils;
import com.fasterxml.jackson.core.type.TypeReference;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.search.*;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.search.suggest.Suggest;
import org.elasticsearch.search.suggest.SuggestBuilder;
import org.elasticsearch.search.suggest.SuggestBuilders;
import org.elasticsearch.search.suggest.SuggestionBuilder;
import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
import org.elasticsearch.search.suggest.completion.CompletionSuggestionBuilder;
import org.elasticsearch.search.suggest.completion.FuzzyOptions;
import org.elasticsearch.search.suggest.phrase.PhraseSuggestionBuilder;
import org.elasticsearch.search.suggest.term.TermSuggestionBuilder;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

@Slf4j
@Service
@Primary
public class ESQueryServiceImpl implements ESQueryService {

    private RestHighLevelClient restHighLevelClient;

    @Resource
    public void setRestHighLevelClient(RestHighLevelClient restHighLevelClient) {
        this.restHighLevelClient = restHighLevelClient;
    }

    /**
     * 查询
     **/
    @Override
    public SearchResponse search(SearchRequest request) {
        try {
            return restHighLevelClient.search(request, RequestOptions.DEFAULT);
        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", request, e.getMessage());
        }
        return null;
    }

    /**
     * 查询
     **/
    @Override
    public <T> List<T> search(SearchSourceBuilder builder, Class<T> c, String... idxNames) {
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        List<T> res = new ArrayList<>();
        try {
            SearchResponse response = search(request);
            SearchHit[] hits = response.getHits().getHits();
            for (SearchHit hit : hits) {
                T t = JsonUtils.jsonToObj(hit.getSourceAsString(), c);
                if (t != null) {
                    res.add(t);
                }
            }
        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", builder, e.getMessage());
        }
        return res;
    }

    /**
     * 查询分页
     **/
    @Override
    public <T> PageResult<T> search(SearchSourceBuilder builder, Class<T> c, int pageNo, int pageSize, String... idxNames) {
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        List<T> res = new ArrayList<>();
        PageResult<T> pagers = new PageResult<>(res);
        try {
            SearchResponse response = search(request);
            long total = response.getHits().getTotalHits().value;
            SearchHit[] hits = response.getHits().getHits();
            for (SearchHit hit : hits) {
                T t = JsonUtils.jsonToObj(hit.getSourceAsString(), c);
                if (t != null) {
                    res.add(t);
                }
            }
            pagers = new PageResult<>((int) total, pageNo, pageSize, false, res);
        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", builder, e.getMessage());
        }
        return pagers;
    }

    /**
     * 查询分页
     **/
    @Override
    public PageResult<Object> search(SearchSourceBuilder builder, int pageNo, int pageSize, String... idxNames) {
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        List<Object> res = new ArrayList<>();
        PageResult<Object> pagers = null;
        try {
            SearchResponse response = search(request);
            long total = response.getHits().getTotalHits().value;
            SearchHit[] hits = response.getHits().getHits();
            for (SearchHit hit : hits) {
                Map<String, Object> map = JsonUtils.jsonToObj(hit.getSourceAsString(), new TypeReference<HashMap<String, Object>>() {
                });
                if (!CollectionUtils.isEmpty(map)) {
                    map.put("score", hit.getScore());
                    res.add(map);
                }
            }
            pagers = new PageResult<>((int) total, pageNo, pageSize, false, res);
        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", builder, e.getMessage());
        }
        return pagers;
    }

    /**
     * 大数据量查询，利用scroll分页
     */
    @Override
    public SearchResponse searchScrollResponse(SearchScrollRequest scrollRequest) {
        try {
            return restHighLevelClient.scroll(scrollRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 大数据量查询，利用scroll分页
     */
    @Override
    public List<ElasticEntity> searchScroll(SearchSourceBuilder builder, String... idxNames) {
        final Scroll scroll = new Scroll(TimeValue.timeValueHours(1));
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        request.scroll(scroll);

        List<ElasticEntity> list = new ArrayList<>();

        try {
            SearchResponse searchResponse = search(request);
            list = getScrollEntity(searchResponse, scroll);

        } catch (IOException e) {
            e.printStackTrace();
            log.error(e.getMessage());
        }
        return list;
    }

    private void searchForResponse(SearchRequestBuilder requestBuilder, Consumer<SearchResponse> consumer) {
        if (requestBuilder == null || consumer == null) {
            return;
        }
        //requestBuilder.setSize(100); //在构建查询条件时，即可设置大小
        //设置排序字段
        SearchRequestBuilder sortBuilder = requestBuilder.addSort("id", SortOrder.ASC);
        SearchResponse searchResponse = sortBuilder.get();
        SearchHit[] searchHits = searchResponse.getHits().getHits();
        while (searchHits.length > 0) {
            consumer.accept(searchResponse);
            SearchHit last = searchHits[searchHits.length - 1];
            sortBuilder = sortBuilder.searchAfter(last.getSortValues());
            searchResponse = sortBuilder.get();
            searchHits = searchResponse.getHits().getHits();
        }
    }

    /**
     * 大数据量查询，利用scroll_after分页
     */
    @Override
    public List<ElasticEntity> searchScrollAfter(SearchSourceBuilder builder, String... idxNames) {
        List<ElasticEntity> elasticEntities = new ArrayList<>();
        Object[] objects = new Object[]{};
        ElasticEntity elasticEntity;
        try {
            SearchRequest searchRequest = new SearchRequest(idxNames);
            builder.size(1000);
            builder.sort("id", SortOrder.ASC);

            while (true) {
                if (objects.length > 0) {
                    builder.searchAfter(objects);
                }
                searchRequest.source(builder);
                SearchResponse searchResponse = search(searchRequest);
                SearchHit[] hits = searchResponse.getHits().getHits();
                if (hits.length == 0) {
                    break;
                }
                objects = hits[hits.length - 1].getSortValues();
                for (SearchHit hit : hits) {
                    elasticEntity = new ElasticEntity();
                    elasticEntity.setId(hit.getId());
                    elasticEntity.setData(hit.getSourceAsMap());
                    elasticEntities.add(elasticEntity);
                }
            }
        } catch (Exception e) {
            log.error(e.getMessage());
        }
        log.info("searchScrollAfter: " + elasticEntities.size());
        return elasticEntities;
    }

    /**
     * 查询
     **/
    @Override
    public SearchHit[] search(SearchSourceBuilder builder, String... idxNames) {
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        try {
            SearchResponse response = search(request);
            return response.getHits().getHits();

        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", builder, e.getMessage());
        }
        return null;
    }

    /**
     * 建议
     *
     * @param idxNames 索引名称
     * @pdaram builder  建议参数配置
     */
    @Override
    public <T> List<T> suggest(SearchSourceBuilder builder, Class<T> c, String... idxNames) {
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        List<T> res = new ArrayList<>();
        try {
            SearchResponse response = search(request);
            Suggest suggestions = response.getSuggest();
            List<T> suggestList = new ArrayList<>();
            for (String key : builder.suggest().getSuggestions().keySet()) {
                suggestList.addAll(getSuggestResults(suggestions, key, c));
            }
            return suggestList;
        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", builder, e.getMessage());
        }
        return res;
    }

    /**
     * 建议
     *
     * @param builder  建议参数配置
     * @param idxNames 索引名称
     */
    @Override
    public Suggest suggest(SearchSourceBuilder builder, String... idxNames) {
        SearchRequest request = new SearchRequest(idxNames);
        request.source(builder);
        try {
            SearchResponse response = search(request);
            return response.getSuggest();
        } catch (Exception e) {
            log.error("search 异常，suggestBuilder={},异常信息={}", builder, e.getMessage());
        }
        return null;
    }

    /**
     * 自动完成
     *
     * @param suggestBuilder 指定在哪个字段搜索以及输入文本等参数配置
     * @param suggestionName 建议名称
     * @param idxNames       索引名称
     */
    @Override
    public <T> List<T> suggest(SuggestionBuilder<?> suggestBuilder, String suggestionName, Class<T> c, String... idxNames) {
        SearchRequest searchRequest = new SearchRequest(idxNames)
                .source(new SearchSourceBuilder()
                        .suggest(new SuggestBuilder()
                                .addSuggestion(suggestionName, suggestBuilder)));

        try {
            SearchResponse searchResponse = search(searchRequest);
            Suggest suggestions = searchResponse.getSuggest();
            return getSuggestResults(suggestions, suggestionName, c);
        } catch (Exception e) {
            log.error(suggestionName + " 异常，suggestBuilder={},异常信息={}", suggestBuilder, e.getMessage());
        }

        return new ArrayList<>();
    }

    /**
     * 词条建议器
     *
     * @param suggestField    建议字段
     * @param suggestMaxCount 返回建议数据大小
     * @param query           输入
     * @param indexNames      索引名称
     **/
    @Override
    public <T> List<T> termSuggest(String suggestField, String query, Integer suggestMaxCount, Class<T> c, String... indexNames) {
        TermSuggestionBuilder termSuggestionBuilder = SuggestBuilders.termSuggestion(suggestField)
                .text(query).size(suggestMaxCount).prefixLength(1).minWordLength(1);

        return suggest(termSuggestionBuilder, "term-suggest", c);
    }

    /**
     * 词组建议器
     *
     * @param suggestField    建议字段
     * @param suggestMaxCount 返回建议数据大小
     * @param query           输入
     * @param indexNames      索引名称
     **/
    @Override
    public <T> List<T> phraseSuggest(String suggestField, String query, Integer suggestMaxCount, Class<T> c, String... indexNames) {
        PhraseSuggestionBuilder phraseSuggestionBuilder = SuggestBuilders.phraseSuggestion(suggestField)
                .text(query).size(suggestMaxCount);//.gramSize(2)

        return suggest(phraseSuggestionBuilder, "phrase-suggest", c);
    }

    /**
     * 自动完成建议器
     *
     * @param suggestField    建议字段
     * @param suggestMaxCount 返回建议数据大小
     * @param query           输入
     * @param indexNames      索引名称
     **/
    @Override
    public <T> List<T> completionSuggest(String suggestField, String query, Integer suggestMaxCount, Class<T> c, String... indexNames) {
        CompletionSuggestionBuilder phraseSuggestionBuilder = SuggestBuilders.completionSuggestion(suggestField)
                .size(suggestMaxCount).prefix(query, getFuzzyOptions(1, 1));

        return suggest(phraseSuggestionBuilder, "completion-suggest", c, indexNames);
    }

    /**
     * 模糊选项
     *
     * @param fuzzyMinLength
     * @param fuzzyPrefixLength
     */
    public FuzzyOptions getFuzzyOptions(int fuzzyMinLength, int fuzzyPrefixLength) {
        return FuzzyOptions.builder()
                .setFuzziness(Fuzziness.AUTO)
                .setFuzzyMinLength(fuzzyMinLength)
                .setUnicodeAware(true)
                .setFuzzyPrefixLength(fuzzyPrefixLength)
                .build();
    }


    /**
     * 解析建议结果成实体集合
     *
     * @param suggestions 建议结果
     * @param suggestName 建议名称
     * @param c           实体类型
     */
    @Override
    public <T> List<T> getSuggestResults(Suggest suggestions, String suggestName, Class<T> c) {
        List<T> entityList = new ArrayList<>();
        if (suggestions != null) {
            Suggest.Suggestion result = suggestions.getSuggestion(suggestName);
            for (Object term : result.getEntries()) {
                if (term instanceof CompletionSuggestion.Entry) {
                    CompletionSuggestion.Entry item = (CompletionSuggestion.Entry) term;
                    if (!item.getOptions().isEmpty()) {
                        for (CompletionSuggestion.Entry.Option option : item.getOptions()) {
                            T t = JsonUtils.jsonToObj(option.getHit().getSourceAsString(), c);
                            if (t != null) {
                                entityList.add(t);
                            }
                        }
                    }
                }
            }
        }
        return entityList;
    }

    /**
     * 解析搜索结果为实体集
     */
    private List<ElasticEntity> getScrollEntity(SearchResponse searchResponse, Scroll scroll) throws IOException {
        List<ElasticEntity> list = new ArrayList<>();
        String scrollId = searchResponse.getScrollId();
        SearchHit[] searchHits = searchResponse.getHits().getHits();
        long totalHits = searchResponse.getHits().getTotalHits().value;
        log.info("es查询集合总数为{}", totalHits);

        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();

        ElasticEntity elasticEntity = null;
        while (searchHits != null && searchHits.length > 0) {
            if (StringUtils.hasLength(scrollId)) {
                clearScrollRequest.addScrollId(scrollId);
            }
            for (SearchHit hit : searchResponse.getHits().getHits()) {
                elasticEntity = new ElasticEntity();
                elasticEntity.setId(hit.getId());
                elasticEntity.setData(hit.getSourceAsMap());
                list.add(elasticEntity);
            }
            searchResponse = searchScrollResponse(new SearchScrollRequest(scrollId).scroll(scroll));
            scrollId = searchResponse.getScrollId();
            searchHits = searchResponse.getHits().getHits();
        }

        if(list.size() > 0) {
            // 清除滚屏
            ClearScrollResponse clearScrollResponse = restHighLevelClient.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
            if (!clearScrollResponse.isSucceeded()) {
                log.error("清除ES游标失败！");
//            throw new RRException("清除es游标失败");
            }
        }

        return list;
    }

    /**
     * 清除滚屏
     */
    @Override
    public void clearScroll(String scrollId) throws IOException {
        if (scrollId == null) {
            return;
        }
        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
        clearScrollRequest.addScrollId(scrollId);
        ClearScrollResponse clearScrollResponse = restHighLevelClient.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
        boolean succeeded = clearScrollResponse.isSucceeded();
    }

    /**
     * 多索引搜索
     *
     * @param searchRequestMap 请求参数集合，key为索引名称，value为SearchSourceBuilder
     * @return
     */
    @Override
    public Map<String, Object> multiSearch(Map<String, SearchSourceBuilder> searchRequestMap) {
        // 存放返回的集合，key为索引名称，value为数据集
        Map<String, Object> dataMap = new HashMap<>();
        if (searchRequestMap == null || searchRequestMap.isEmpty()) {
            return dataMap;
        }

        // 创建搜索请求体
        for (Map.Entry<String, SearchSourceBuilder> entry : searchRequestMap.entrySet()) {
            String index = entry.getKey();
            SearchSourceBuilder builder = entry.getValue();
            List<Map<String, Object>> dataList = new ArrayList<>();
            try {
                dataList = batchSearch(builder, index);
            } catch (Exception e) {
                log.error("multi search 异常，searchSourceBuilder={},异常信息={}", builder, e.getMessage());
            }
            dataMap.put(index, dataList);
        }
        return dataMap;
    }

    /**
     * 大数据量，分批次搜索
     *
     * @param builder
     * @param indexName
     * @return
     */
    @SuppressWarnings("unchecked")
    private List<Map<String, Object>> batchSearch(SearchSourceBuilder builder, String indexName) {
        SearchRequest searchRequest = new SearchRequest(indexName);
        int count = builder.size();
        // 批量查询条数
        int batchCount = 1000;
        // 存储搜索完成的数据集
        List<Map<String, Object>> resultList = new ArrayList<>();
        // 存储每次搜索最后一条数据的位置
        Object[] sortValue = null;
        // 存储每次搜索的数据集
        List<Map<String, Object>> searchList = new ArrayList<>();
        if (count == 0) {
            count = batchCount;
        }
        // 遍历计数
        int i = 0;
        // 遍历总次数
        int num = count / batchCount;
        while (resultList.size() < count && i <= num) {
            int realCount;
            if (i == num) {
                // 找出最后一次遍历的数据量
                realCount = count % batchCount;
                if (realCount == 0) {
                    realCount = batchCount;
                }
            } else {
                realCount = batchCount;
            }

            if (sortValue != null) {
                builder.searchAfter(sortValue);
                // 使用search_after参数，from参数必须设置为0
                builder.from(0);
            }

            builder.size(realCount);
            searchRequest.source(builder);
            // 请求es搜索
            Map<String, Object> dataMap = searchAsync(searchRequest);
            sortValue = (Object[]) dataMap.get("sortValue");
            searchList = (List<Map<String, Object>>) dataMap.get("dataList");
            if (!searchList.isEmpty()) {
                resultList.addAll(searchList);
            }

            i++;
        }

        return resultList;
    }

    /**
     * 异步搜索，返回结果为数据集和最后一条数据标志位（可用于滚动搜索）
     *
     * @param searchRequest
     * @return
     */
    private Map<String, Object> searchAsync(SearchRequest searchRequest) {
        // 存储数据集和最后一条数据标志位
        CompletableFuture<Map<String, Object>> future = new CompletableFuture<>();
        restHighLevelClient.searchAsync(searchRequest, RequestOptions.DEFAULT, new ActionListener<SearchResponse>() {
            @Override
            public void onResponse(SearchResponse searchResponse) {
                // 存储搜索数据的集合
                List<Map<String, Object>> dataList;
                // 存储分页标志位
                Object[] sortValue = null;
                SearchHits searchHits = searchResponse.getHits();
                SearchHit[] hits = searchHits.getHits();

                // 获取高亮显示参数
                SearchSourceBuilder builder = searchRequest.source();
                HighlightBuilder highlightBuilder = builder.highlighter();
                // 若有高亮配置，则遍历数据集添加高亮结果
                if (highlightBuilder != null) {
                    dataList = getResultWithHighlight(hits);
                } else {
                    dataList = Stream.of(hits).map(SearchHit::getSourceAsMap).collect(Collectors.toList());
                }
                // 获取最后一条数据的位置
                if (hits.length > 1) {
                    sortValue = searchHits.getAt(hits.length - 1).getSortValues();
                }

                Map<String, Object> result = new HashMap<>();
                result.put("sortValue", sortValue);
                result.put("dataList", dataList);
                future.complete(result);
            }

            @Override
            public void onFailure(Exception e) {
                future.completeExceptionally(e);
            }
        });
        return future.join();
    }

    /**
     * 往数据集添加高亮结果
     *
     * @param hits
     * @return
     */
    private List<Map<String, Object>> getResultWithHighlight(SearchHit[] hits) {
        List<Map<String, Object>> dataList = Stream.of(hits).map(hit -> {
            Map<String, Object> hitMap = hit.getSourceAsMap();
            Map<String, HighlightField> hlFields = hit.getHighlightFields();
            if (hitMap != null && hlFields != null && !hlFields.isEmpty()) {
                // 存储每条数据高亮结果的map
                Map<String, Object> hlMap = new HashMap<>();
                hlFields.forEach((key, field) -> {
                    if (field != null) {
                        String name = field.getName();
                        Text[] fragments = field.fragments();
                        hlMap.put(name, Arrays.stream(fragments).map(Text::string).collect(Collectors.toList()));
                    }
                });
                hitMap.put("highlight", hlMap);
            }
            return hitMap;
        }).collect(Collectors.toList());
        return dataList;
    }
}
